﻿#include "onnxdetector.h"
#include <fstream>
#include <iomanip>
#include <algorithm>
#include <cmath>

OnnxDetector::OnnxDetector(const std::string& _modelPath, const std::vector<std::string>& class_names, float confidenceThreshold, float nmsThreshold)
: modelPath(_modelPath), CLASS_NAMES(class_names),  CONFIDENCE_THRESHOLD(confidenceThreshold), IOU_THRESHOLD(nmsThreshold)
{
    init();
}

void OnnxDetector::init(){
    try {
        env_ = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "YOLOv8_Detector");
        // 配置会话选项
        session_options_.SetIntraOpNumThreads(1);
        session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

        // 启用 CUDA 加速（可选）
        // OrtSessionOptionsAppendExecutionProvider_CUDA(session_options_, 0);

        // 加载模型，跨平台模型路径处理
        std::string model_path_str = modelPath;
        #ifdef _WIN32
            std::wstring w_model_path = std::wstring(model_path_str.begin(), model_path_str.end());
            const ORTCHAR_T* model_path_ort = w_model_path.c_str();
        #else
            const ORTCHAR_T* model_path_ort = model_path_str.c_str();
        #endif
        // 创建会话
        session_ = std::make_unique<Ort::Session>(*env_, w_model_path.c_str(), session_options_);

    }
    catch (const Ort::Exception& e) {
        throw;
    }
}

// 图像预处理
std::tuple<cv::Mat, OnnxDetector::TransformParams> OnnxDetector::preprocess_image(const cv::Mat& image) {
    // 计算缩放比例并保持宽高比
    float scale = std::min(static_cast<float>(INPUT_SIZE) / image.cols,
                          static_cast<float>(INPUT_SIZE) / image.rows);
    int scaled_width = static_cast<int>(image.cols * scale);
    int scaled_height = static_cast<int>(image.rows * scale);

    // 计算填充偏移量
    int offset_x = (INPUT_SIZE - scaled_width) / 2;
    int offset_y = (INPUT_SIZE - scaled_height) / 2;

    // 创建目标图像(640x640)并填充灰色边框
    cv::Mat scaled_image(INPUT_SIZE, INPUT_SIZE, CV_8UC3, cv::Scalar(128, 128, 128));

    // 缩放原始图像并放置到中央
    cv::Mat resized;
    cv::resize(image, resized, cv::Size(scaled_width, scaled_height), 0, 0, cv::INTER_LINEAR);
    resized.copyTo(scaled_image(cv::Rect(offset_x, offset_y, scaled_width, scaled_height)));

    // 转换到浮点型并归一化
    cv::Mat float_image;
    scaled_image.convertTo(float_image, CV_32FC3, 1.0 / 255.0);

    // 转换为 CHW 格式 [1, 3, H, W]
    cv::Mat chw_image = cv::dnn::blobFromImage(float_image);

    TransformParams params;
    params.scale = scale;
    params.offset_x = offset_x;
    params.offset_y = offset_y;
    params.scaled_width = scaled_width;
    params.scaled_height = scaled_height;

    return {chw_image, params};
}

// 解析模型输出
std::vector<OnnxDetector::Detection> OnnxDetector::parse_model_output(const cv::Mat& input_tensor, const TransformParams& params, const cv::Mat& original_image) {
    // 获取输入输出信息
    Ort::AllocatorWithDefaultOptions allocator;
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

    // 准备输入
    std::vector<int64_t> input_shape = {1, 3, INPUT_SIZE, INPUT_SIZE};
    std::vector<Ort::Value> input_tensors;
    input_tensors.push_back(Ort::Value::CreateTensor<float>(
        memory_info,
        (float*)input_tensor.data,
        input_tensor.total(),
        input_shape.data(),
        input_shape.size()
    ));

    // 输入节点名称
    std::vector<const char*> input_node_names = {"images"};

    // 输出节点名称
    std::vector<const char*> output_node_names = {"output0"};

    // 运行推理
    auto output_tensors = session_->Run(
        Ort::RunOptions{nullptr},
        input_node_names.data(),
        input_tensors.data(),
        input_tensors.size(),
        output_node_names.data(),
        output_node_names.size()
    );

    // 获取输出张量
    float* output_data = output_tensors[0].GetTensorMutableData<float>();
    auto output_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();

    // 输出形状: [batch, 84, 8400]
    int num_classes = CLASS_NAMES.size();
    int num_boxes = output_shape[2];

    // 解析检测结果
    std::vector<Detection> detections;

    for (int i = 0; i < num_boxes; ++i) {
        // 获取边界框数据
        float x_center = output_data[i];
        float y_center = output_data[num_boxes + i];
        float width = output_data[2 * num_boxes + i];
        float height = output_data[3 * num_boxes + i];

        // 跳过无效边界框
        if (width <= 0 || height <= 0) continue;

        // 找到最大置信度的类别
        float max_confidence = 0.0f;
        int class_id = -1;

        for (int c = 0; c < num_classes; ++c) {
            float confidence = output_data[(4 + c) * num_boxes + i];
            if (confidence > max_confidence) {
                max_confidence = confidence;
                class_id = c;
            }
        }

        // 过滤低置信度的检测结果
        if (max_confidence < CONFIDENCE_THRESHOLD) continue;

        // 将中心坐标转换为角坐标
        float x1 = x_center - width / 2;
        float y1 = y_center - height / 2;
        float x2 = x_center + width / 2;
        float y2 = y_center + height / 2;

        // 关键修复：将坐标从640x640空间转换回原始图像空间
        // 1. 减去填充偏移量
        x1 = (x1 - params.offset_x) / params.scale;
        y1 = (y1 - params.offset_y) / params.scale;
        x2 = (x2 - params.offset_x) / params.scale;
        y2 = (y2 - params.offset_y) / params.scale;

        // 2. 限制在原始图像范围内
        x1 = std::max(0.0f, std::min(x1, static_cast<float>(original_image.cols)));
        y1 = std::max(0.0f, std::min(y1, static_cast<float>(original_image.rows)));
        x2 = std::max(0.0f, std::min(x2, static_cast<float>(original_image.cols)));
        y2 = std::max(0.0f, std::min(y2, static_cast<float>(original_image.rows)));

        // 创建检测结果
        Detection detection;
        detection.bbox = cv::Rect_<float>(x1, y1, x2 - x1, y2 - y1);
        detection.confidence = max_confidence;
        detection.class_id = class_id;
        detection.label = CLASS_NAMES[class_id];

        detections.push_back(detection);
    }

    return detections;
}

// 计算两个矩形的IoU(交并比)
float OnnxDetector::calculate_iou(const cv::Rect_<float>& rect1, const cv::Rect_<float>& rect2) {
    float x1 = std::max(rect1.x, rect2.x);
    float y1 = std::max(rect1.y, rect2.y);
    float x2 = std::min(rect1.x + rect1.width, rect2.x + rect2.width);
    float y2 = std::min(rect1.y + rect1.height, rect2.y + rect2.height);

    float intersection = std::max(0.0f, x2 - x1) * std::max(0.0f, y2 - y1);
    float area1 = rect1.width * rect1.height;
    float area2 = rect2.width * rect2.height;

    return intersection / (area1 + area2 - intersection);
}

// 非极大值抑制(NMS)算法
std::vector<OnnxDetector::Detection> OnnxDetector::apply_nms(std::vector<Detection>& detections) {
    // 按置信度降序排序
    std::sort(detections.begin(), detections.end(), [](const Detection& a, const Detection& b) {
        return a.confidence > b.confidence;
    });

    std::vector<Detection> final_detections;

    while (!detections.empty()) {
        // 取当前最高置信度的检测结果
        Detection current = detections[0];
        final_detections.push_back(current);
        detections.erase(detections.begin());

        // 移除与当前框重叠度高的框
        for (auto it = detections.begin(); it != detections.end();) {
            float iou = calculate_iou(current.bbox, it->bbox);
            if (iou > IOU_THRESHOLD) {
                it = detections.erase(it);
            } else {
                ++it;
            }
        }
    }

    return final_detections;
}

// 主检测函数
std::vector<OnnxDetector::Detection> OnnxDetector::detect(const cv::Mat& image) {
    // 预处理图像
    auto [input_tensor, transform_params] = preprocess_image(image);

    // 解析模型输出
    auto detections = parse_model_output(input_tensor, transform_params, image);

    // 应用非极大值抑制
    return apply_nms(detections);
}

// 为每个类别生成颜色
cv::Scalar OnnxDetector::get_class_color(int class_id) {
    // 使用 HSV 颜色空间确保颜色区分度高
    float hue = (class_id * 137.5f); // 黄金角间隔
    int hi = static_cast<int>(hue / 60) % 6;
    float f = hue / 60 - static_cast<int>(hue / 60);

    float v = 0.9f * 255;
    float p = v * (1 - 0.9f);
    float q = v * (1 - f * 0.9f);
    float t = v * (1 - (1 - f) * 0.9f);

    switch (hi) {
        case 0: return cv::Scalar(t, p, v); // BGR
        case 1: return cv::Scalar(q, v, p);
        case 2: return cv::Scalar(p, v, t);
        case 3: return cv::Scalar(p, q, v);
        case 4: return cv::Scalar(t, p, v);
        default: return cv::Scalar(v, p, q);
    }
}

// 在图像上绘制检测结果
cv::Mat OnnxDetector::draw_detections(const cv::Mat& image, const std::vector<Detection>& detections) {
    cv::Mat result_image = image.clone();
    int thickness = std::max(1, static_cast<int>(image.cols / 400));

    for (const auto& detection : detections) {
        // 获取类别颜色
        cv::Scalar color = get_class_color(detection.class_id);

        // 绘制边界框
        cv::rectangle(result_image, detection.bbox, color, thickness);

        // 创建标签文本
        std::ostringstream label_ss;
        label_ss << detection.label << ": " << std::fixed << std::setprecision(2) << detection.confidence;
        std::string label = label_ss.str();

        // 计算文本大小
        int baseline = 0;
        int font_face = cv::FONT_HERSHEY_SIMPLEX;
        double font_scale = 0.5;
        cv::Size text_size = cv::getTextSize(label, font_face, font_scale, 1, &baseline);

        // 计算标签背景位置
        int text_x = static_cast<int>(detection.bbox.x);
        int text_y = static_cast<int>(detection.bbox.y) - 5;
        text_y = std::max(0, text_y);

        // 绘制标签背景
        cv::rectangle(
            result_image,
            cv::Point(text_x, text_y - text_size.height - 5),
            cv::Point(text_x + text_size.width, text_y + 5),
            color,
            cv::FILLED
        );

        // 绘制标签文本
        cv::putText(
            result_image,
            label,
            cv::Point(text_x, text_y - 3),
            font_face,
            font_scale,
            cv::Scalar(255, 255, 255),
            1
        );
    }

    return result_image;
}

// 保存检测结果到文本文件
void OnnxDetector::save_detections_to_file(const std::vector<Detection>& detections, const std::string& filename) {
    std::ofstream out_file(filename);
    if (!out_file) {
        std::cerr << "无法打开输出文件: " << filename << std::endl;
        return;
    }

    out_file << "检测结果:\n";
    out_file << "类别\t置信度\tX\tY\t宽度\t高度\n";

    for (const auto& detection : detections) {
        out_file << detection.label << "\t"
                 << std::fixed << std::setprecision(4) << detection.confidence << "\t"
                 << static_cast<int>(detection.bbox.x) << "\t"
                 << static_cast<int>(detection.bbox.y) << "\t"
                 << static_cast<int>(detection.bbox.width) << "\t"
                 << static_cast<int>(detection.bbox.height) << "\n";
    }

    out_file.close();
}

OnnxDetector::~OnnxDetector() {

}
